import argparse
import zipfile
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.asymmetric import padding
from cryptography.hazmat.primitives.asymmetric import rsa
import os


def rmdir(path):
    if os.path.isdir(path):
        for file in os.listdir(path):
            rmdir(os.path.join(path, file))
        if os.path.exists(path):
            os.rmdir(path)
    else:
        if os.path.exists(path):
            os.remove(path)


class FileCrypto(object):
    def __init__(self, path):
        self.path = path
        self.decrypt_file = []
        self.key_file = None

    def encrypt(self, pwd=None):
        father_path = os.path.dirname(self.path)
        zip_name = f"{os.path.basename(self.path)}.zip"
        zip_full_path = os.path.join(father_path, zip_name)
        z = zipfile.ZipFile(zip_full_path, 'w')

        try:

            private_key = rsa.generate_private_key(
                public_exponent=65537,
                key_size=2048,
            )
            public_key = private_key.public_key()
            pem = private_key.private_bytes(
                encoding=serialization.Encoding.PEM,
                format=serialization.PrivateFormat.PKCS8,
                encryption_algorithm=serialization.BestAvailableEncryption(pwd.encode('utf-8'))
            )

            z.writestr(zinfo_or_arcname='key.pem', data=pem)
            self.recursive_zip_encrypto(self.path, z, self.path, public_key)
            z.close()
            rmdir(self.path)
        except Exception as e:
            print(e)
            rmdir(zip_full_path)


    def decrypt(self, pwd=None, dest_path=None):
        if not dest_path:
            base_name = os.path.basename(os.path.splitext(self.path)[0])
            dest_path = os.path.join(os.path.dirname(self.path), base_name)
        try:
            self.un_zip(self.path, dest_path)
            with open(self.key_file, "rb") as key_file:
                private_key = serialization.load_pem_private_key(
                    key_file.read(),
                    password=pwd.encode('utf-8'),
                )

            for df in self.decrypt_file:
                plaintext = b''
                with open(df, 'rb') as fr:
                    while True:
                        content_byte = fr.read(256)
                        if not len(content_byte):
                            break
                        plaintext += private_key.decrypt(
                            content_byte,
                            padding.OAEP(
                                mgf=padding.MGF1(algorithm=hashes.SHA256()),
                                algorithm=hashes.SHA256(),
                                label=None
                            )
                        )

                with open(df, 'wb') as fw:
                    fw.write(plaintext)
            os.remove(self.key_file)
            rmdir(self.path)

        except Exception as e:
            print(e)
            rmdir(dest_path)

    @classmethod
    def recursive_zip_encrypto(cls, zip_path, zip_obj, zip_root_path, pk=None):
        if os.path.isdir(zip_path):
            for i in os.listdir(zip_path):
                cls.recursive_zip_encrypto(os.path.join(zip_path, i), zip_obj, zip_root_path, pk)
        else:
            if zip_path == zip_root_path:
                zip_name = os.path.basename(zip_path)
            else:
                zip_name = os.path.relpath(zip_path, zip_root_path)
            if pk:
                ciphertext = b''
                with open(zip_path, 'rb') as fr:
                    while True:
                        content = fr.read(128)
                        if not len(content):
                            break
                        ciphertext += pk.encrypt(
                            content,
                            padding.OAEP(
                                mgf=padding.MGF1(algorithm=hashes.SHA256()),
                                algorithm=hashes.SHA256(),
                                label=None
                            )
                        )

                zip_obj.writestr(zinfo_or_arcname=zip_name, data=ciphertext)
            else:
                zip_obj.write(zip_path, zip_name)

    def un_zip(self, file_name, dest):
        """unzip zip file"""
        zip_file = zipfile.ZipFile(file_name)
        for names in zip_file.namelist():
            a = zip_file.extract(names, path=dest)
            if os.path.basename(a) == 'key.pem':
                self.key_file = a
            else:
                self.decrypt_file.append(a)
        zip_file.close()

    @classmethod
    def recursive_zip(cls, zip_path, zip_obj, zip_root_path):
        if os.path.isdir(zip_path):
            for i in os.listdir(zip_path):
                cls.recursive_zip(os.path.join(zip_path, i), zip_obj, zip_root_path)
        else:
            if zip_path == zip_root_path:
                zip_name = os.path.basename(zip_path)
            else:
                zip_name = os.path.relpath(zip_path, zip_root_path)
            zip_obj.write(zip_path, zip_name)


def get_arg_parser():
    parser = argparse.ArgumentParser(
        description=(
            'Enter a password to encrypt the compressed folder '
        ),
        add_help=False
    )
    parser.add_argument('PATH', help='file path')
    parser.add_argument(
        '-h',
        '--help',
        action='help',
        help='ceason'
    )

    parser.add_argument(
        '-m',
        action='store',
        dest='MODE',
        choices=['encrypt', 'decrypt'],
        type=str,
        help='mode'
    )

    parser.add_argument(
        '-p',
        action='store',
        dest='PASSWORD',
        default=None,
        type=str,
        help='password'
    )
    return parser


if __name__ == '__main__':
    args = get_arg_parser().parse_args()

    if not os.path.exists(args.PATH):
        print("file not found")
    else:
        fc = FileCrypto(args.PATH)

        if args.MODE == "encrypt":
            fc.encrypt(args.PASSWORD)
        elif args.MODE == "decrypt":
            fc.decrypt(args.PASSWORD)
        else:
            pass
